"""
@Description :   加载 Obj365 数据集
@Author      :   tqychy 
@Time        :   2025/08/27 20:13:12
"""

import json
import os

from torch.utils.data import Dataset


class Obj365Dataset(Dataset):
    def __init__(self, *args, data_root="./dataset/raw/obj365/val", **kwargs):

        super().__init__()
        self.cfg, self.logger = args
        self.data_root = data_root
        raw_data = json.load(open(os.path.join(data_root, "zhiyuan_objv2_val.json"), "r", encoding="utf-8"))
        self.indices, self.data = self.pre_process(raw_data)
        max_len = kwargs.get("max_len", -1)
        self.max_len = -1 if max_len == -1 else min(max_len, len(self.indices))

    def pre_process(self, raw_data):
        """
        将 raw 整理成包含图片路径、类别和 bbox 的字典的格式
        """
        # 整理类别
        category_dict = {}
        for cat in raw_data["categories"]:
            category_dict[cat["id"]] = cat["name"]
        
        data = {}
        for image in raw_data["images"]:
            image_id = image["id"]
            image_path = image["file_name"]
            path, image_name = os.path.split(image_path)
            image_path = os.path.join(self.data_root, os.path.basename(path), image_name)

            data[image_id] = {
                "image_path": image_path,
                "categories": [],
                "bboxes": []
            }
        for ann in raw_data["annotations"]:
            image_id = ann["image_id"]
            category_id = ann["category_id"]
            bbox = ann["bbox"]
            data[image_id]["categories"].append(category_dict[category_id])
            data[image_id]["bboxes"].append(bbox)

        return list(data.keys()), data
    
    def __len__(self):
        return len(self.indices) if self.max_len == -1 else self.max_len

    def __getitem__(self, idx):
        image_id = self.indices[idx]
        sample = self.data[image_id]

        image_path = sample["image_path"]
        categories = sample["categories"]
        bboxes = sample["bboxes"]

        return {
            "image_path": image_path,
            "bbox": bboxes,
            "sentences":categories,
            "category": categories
        }

if __name__ == "__main__":
    import matplotlib
    import matplotlib.pyplot as plt
    import torchvision.transforms as T
    from matplotlib.patches import Rectangle
    from torch.utils.data import DataLoader
    from PIL import Image

    matplotlib.use('Agg')
    save_path = "./dataset/refcoco_sample_test"
    os.makedirs(save_path, exist_ok=True)

    dataset = Obj365Dataset(None, None, max_len=100)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=False)

    print(f"Total {len(dataloader)} samples.")
    for i, data in enumerate(dataloader):
        image_path = data["image_path"][0]
        bboxes = data["bbox"]
        categories = data["sentences"]

        transforms = T.ToTensor()
        image = Image.open(image_path).convert("RGB")
        image = transforms(image)
        image = image.squeeze().permute(1, 2, 0)

        ax = plt.gca()
        ax.imshow(image)
        for category, bbox in zip(categories, bboxes):
            box_plot = Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], fill=False, edgecolor='green', linewidth=3)
            ax.add_patch(box_plot)
            ax.text(bbox[0], bbox[1] - 10, category, fontsize=8, color="purple", bbox=dict(facecolor='white', alpha=0.7, pad=2))

        plt.savefig(os.path.join(save_path, f"{i}.png"))
        plt.clf()
        c = input("Press Enter to continue, q to quit.")
        if c == 'q':
            break  

    # import tarfile
    # def extract_tar_gz_files(directory):
    #     # 遍历目录下的所有文件
    #     for filename in os.listdir(directory):
    #         if filename.endswith('.tar.gz'):
    #             file_path = os.path.join(directory, filename)
    #             print(f"正在解压: {file_path}")
                
    #             try:
    #                 # 打开 tar.gz 文件
    #                 with tarfile.open(file_path, 'r:gz') as tar:
    #                     # 解压到当前目录
    #                     tar.extractall(path=directory)
    #                 print(f"成功解压: {filename}")
    #             except Exception as e:
    #                 print(f"解压 {filename} 时出错: {str(e)}")
    # extract_tar_gz_files("./dataset/raw/obj365/val")